import torch
import torch.nn as nn
import torchvision.models as models


class CNN(nn.Module):
    def __init__(self, class_num):
        super(CNN, self).__init__()
        vgg16 = models.vgg16(pretrained=True)
        conv_blocks = list(vgg16.features.children())[:31]
        self.vgg= nn.Sequential(*conv_blocks)
        self.linear = nn.Linear(25088, class_num)

    def forward(self, x):
        features = self.vgg(x)
        out = features.view(features.size(0), -1)
        output = self.linear(out)
        return output


